import sys
import os
import cv2
import torch
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QPushButton, QLabel,
    QFileDialog, QMessageBox, QVBoxLayout, QHBoxLayout, QWidget
)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from pathlib import Path

# 导入 YOLOv5 检测脚本
from detect_images import run, parse_opt


class DetectionWorker(QThread):
    finished = pyqtSignal(str)  # 检测结果图片路径
    error = pyqtSignal(str)     # 错误信息

    def __init__(self, source_path):
        super().__init__()
        self.source_path = source_path

    def run(self):
        try:
            # 设置 YOLOv5 参数
            opt = parse_opt()
            opt.weights = 'yolov5s.pt'  # 假设模型已下载
            opt.source = self.source_path
            opt.project = 'runs/detect/gui_temp'
            opt.name = 'temp_run'
            opt.exist_ok = True
            opt.nosave = False

            # 运行检测
            run(**vars(opt))

            # 获取保存目录和检测结果图片路径
            save_dir = Path(opt.project) / opt.name
            image_name = Path(self.source_path).name
            detected_image_path = str(save_dir / image_name)

            self.finished.emit(detected_image_path)
        except Exception as e:
            self.error.emit(str(e))


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("YOLOv5 Detection GUI")
        self.setGeometry(100, 100, 1200, 600)
        self.init_ui()
        self.image_path = None

    def init_ui(self):
        # 按钮
        self.select_img_btn = QPushButton("选择图片")
        self.select_folder_btn = QPushButton("选择文件夹")
        self.run_btn = QPushButton("开始检测")

        # 图像展示区域
        self.original_label = QLabel("原图显示区域")
        self.original_label.setAlignment(Qt.AlignCenter)
        self.detected_label = QLabel("检测结果")
        self.detected_label.setAlignment(Qt.AlignCenter)

        # 布局
        btn_layout = QHBoxLayout()
        btn_layout.addWidget(self.select_img_btn)
        btn_layout.addWidget(self.select_folder_btn)
        btn_layout.addWidget(self.run_btn)

        image_layout = QHBoxLayout()
        image_layout.addWidget(self.original_label)
        image_layout.addWidget(self.detected_label)

        main_layout = QVBoxLayout()
        main_layout.addLayout(btn_layout)
        main_layout.addLayout(image_layout)

        container = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

        # 信号槽连接
        self.select_img_btn.clicked.connect(self.select_image)
        self.select_folder_btn.clicked.connect(self.select_folder)
        self.run_btn.clicked.connect(self.run_detection)

    def select_image(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(
            self, "选择图片", "", "Images (*.png *.jpg *.bmp)"
        )
        if file_path:
            self.image_path = file_path
            self.show_image(self.original_label, file_path)

    def select_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder_path:
            self.image_path = folder_path
            # 可选：显示文件夹中的第一张图片
            self.show_first_image_in_folder(folder_path)

    def show_first_image_in_folder(self, folder_path):
        for ext in ('*.png', '*.jpg', '*.bmp'):
            files = list(Path(folder_path).glob(ext))
            if files:
                self.show_image(self.original_label, str(files[0]))
                break

    def run_detection(self):
        if not self.image_path:
            QMessageBox.warning(self, "错误", "请选择图片或文件夹")
            return

        self.worker = DetectionWorker(self.image_path)
        self.worker.finished.connect(self.on_detection_finished)
        self.worker.error.connect(self.on_detection_error)
        self.worker.start()

    def on_detection_finished(self, detected_image_path):
        if not Path(detected_image_path).exists():
            # QMessageBox.warning(self, "错误", "未找到检测结果文件")
            return
        self.show_image(self.detected_label, detected_image_path)
        QMessageBox.information(self, "完成", "检测已完成！")

    def on_detection_error(self, error_msg):
        QMessageBox.critical(self, "错误", f"检测失败: {error_msg}")

    def show_image(self, label, image_path):
        image = cv2.imread(image_path)
        if image is None:
            label.setText("无法加载图片")
            return
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        h, w, ch = image.shape
        bytes_per_line = ch * w
        q_img = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_img)
        label.setPixmap(pixmap.scaled(label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())